#!/usr/bin/env python3
"""
S3 Multipart Upload Tool - MultipartUpload API 演示工具

这个工具演示如何使用 S3 多部分上传 API 进行大文件上传。

使用示例:
    python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file
    python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file --part-size 10485760
    python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --upload-id existing-upload-id --action list-parts
    python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --upload-id existing-upload-id --action abort
"""

import argparse
import boto3
import json
import sys
import os
import hashlib
from datetime import datetime, timezone
from botocore.config import Config
from botocore.exceptions import ClientError, NoCredentialsError


class MultipartUploadClient:
    """S3 多部分上传客户端"""
    
    def __init__(self, endpoint_url='http://localhost:8000', 
                 access_key='AKIAIOSFODNN7EXAMPLE', 
                 secret_key='wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                 region='us-east-1'):
        """
        初始化多部分上传客户端
        
        Args:
            endpoint_url: S3 服务端点
            access_key: AWS Access Key
            secret_key: AWS Secret Key
            region: AWS 区域
        """
        self.endpoint_url = endpoint_url
        self.access_key = access_key
        self.secret_key = secret_key
        self.region = region
        self._client = None
    
    def _get_client(self):
        """获取 S3 客户端"""
        if self._client is None:
            self._client = boto3.client(
                's3',
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
                config=Config(s3={'addressing_style': 'path'})
            )
        return self._client
    
    def initiate_multipart_upload(self, bucket, key):
        """初始化多部分上传"""
        try:
            s3_client = self._get_client()
            print(f"🚀 初始化多部分上传 s3://{bucket}/{key}")
            
            response = s3_client.create_multipart_upload(Bucket=bucket, Key=key)
            upload_id = response['UploadId']
            
            print(f"✅ 初始化成功！UploadId: {upload_id}")
            return upload_id
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 初始化失败 [{error_code}]: {error_message}")
            return None
        except Exception as e:
            print(f"❌ 未知错误: {e}")
            return None
    
    def upload_part(self, bucket, key, upload_id, part_number, data, verify_md5=False):
        """上传分片"""
        try:
            upload_args = {
                'Bucket': bucket,
                'Key': key,
                'PartNumber': part_number,
                'UploadId': upload_id,
                'Body': data
            }
            
            if verify_md5:
                md5_hash = hashlib.md5(data).hexdigest()
                upload_args['ContentMD5'] = md5_hash
                print(f"  📋 Content-MD5: {md5_hash}")
            
            print(f"📤 上传分片 {part_number} ({len(data)} 字节)")
            response = self._get_client().upload_part(**upload_args)
            etag = response['ETag']
            
            print(f"  ✅ 分片 {part_number} 上传成功！ETag: {etag}")
            return {'PartNumber': part_number, 'ETag': etag}
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"  ❌ 分片 {part_number} 上传失败 [{error_code}]: {error_message}")
            return None
        except Exception as e:
            print(f"  ❌ 分片 {part_number} 上传失败: {e}")
            return None
    
    def complete_multipart_upload(self, bucket, key, upload_id, parts):
        """完成多部分上传"""
        try:
            print(f"🔗 完成多部分上传，合并 {len(parts)} 个分片")
            
            response = self._get_client().complete_multipart_upload(
                Bucket=bucket,
                Key=key,
                UploadId=upload_id,
                MultipartUpload={'Parts': parts}
            )
            
            etag = response.get('ETag')
            location = response.get('Location', f"s3://{bucket}/{key}")
            
            print(f"✅ 多部分上传完成！")
            print(f"   📍 位置: {location}")
            print(f"   🏷️  最终 ETag: {etag}")
            return True
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 完成上传失败 [{error_code}]: {error_message}")
            return False
        except Exception as e:
            print(f"❌ 完成上传失败: {e}")
            return False
    
    def abort_multipart_upload(self, bucket, key, upload_id):
        """中止多部分上传"""
        try:
            print(f"🛑 中止多部分上传 {upload_id}")
            
            self._get_client().abort_multipart_upload(
                Bucket=bucket, 
                Key=key, 
                UploadId=upload_id
            )
            
            print("✅ 多部分上传已中止")
            return True
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 中止上传失败 [{error_code}]: {error_message}")
            return False
        except Exception as e:
            print(f"❌ 中止上传失败: {e}")
            return False
    
    def list_parts(self, bucket, key, upload_id):
        """列举分片"""
        try:
            print(f"📋 列举分片 {upload_id}")
            
            response = self._get_client().list_parts(
                Bucket=bucket, 
                Key=key, 
                UploadId=upload_id
            )
            
            parts = response.get('Parts', [])
            print(f"找到 {len(parts)} 个分片:")
            print("-" * 60)
            
            if not parts:
                print("📭 暂无分片")
                return []
            
            total_size = 0
            for part in parts:
                size = part['Size']
                total_size += size
                last_modified = part['LastModified'].strftime('%Y-%m-%d %H:%M:%S')
                print(f"  📄 分片 {part['PartNumber']:>3}: {size:>10} 字节, ETag: {part['ETag']}, 修改时间: {last_modified}")
            
            print("-" * 60)
            print(f"📊 总计: {len(parts)} 个分片, {total_size} 字节 ({total_size / 1024 / 1024:.2f} MB)")
            return parts
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message')
            print(f"❌ 列举分片失败 [{error_code}]: {error_message}")
            return []
        except Exception as e:
            print(f"❌ 列举分片失败: {e}")
            return []
    
    def upload_file(self, bucket, key, file_path, part_size=5*1024*1024, verify_md5=False, upload_id=None):
        """上传文件（多部分上传）"""
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"❌ 错误: 文件不存在: {file_path}")
                return False
            
            file_size = os.path.getsize(file_path)
            print(f"📁 文件信息:")
            print(f"   路径: {file_path}")
            print(f"   大小: {file_size} 字节 ({file_size / 1024 / 1024:.2f} MB)")
            print(f"   分片大小: {part_size} 字节 ({part_size / 1024 / 1024:.2f} MB)")
            
            # 计算分片数量
            part_count = (file_size + part_size - 1) // part_size
            print(f"   预计分片数量: {part_count}")
            
            if part_count > 10000:
                print("❌ 错误: 分片数量超过 10000 个限制")
                return False
            
            # 初始化或使用现有上传
            if upload_id:
                print(f"🔄 使用现有 UploadId: {upload_id}")
            else:
                upload_id = self.initiate_multipart_upload(bucket, key)
                if not upload_id:
                    return False
            
            # 上传分片
            parts = []
            print(f"\n📤 开始上传分片...")
            print("-" * 50)
            
            with open(file_path, 'rb') as f:
                for part_num in range(1, part_count + 1):
                    data = f.read(part_size)
                    if not data:
                        break
                    
                    part_info = self.upload_part(bucket, key, upload_id, part_num, data, verify_md5)
                    if not part_info:
                        print("❌ 上传失败，中止多部分上传")
                        self.abort_multipart_upload(bucket, key, upload_id)
                        return False
                    
                    parts.append(part_info)
            
            # 完成上传
            print(f"\n🔗 完成上传...")
            if not self.complete_multipart_upload(bucket, key, upload_id, parts):
                print("❌ 完成上传失败，中止多部分上传")
                self.abort_multipart_upload(bucket, key, upload_id)
                return False
            
            return True
            
        except Exception as e:
            print(f"❌ 上传失败: {e}")
            return False


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='S3 多部分上传 API 演示工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 上传文件
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file
  
  # 自定义分片大小
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file --part-size 10485760
  
  # 启用 MD5 校验
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file --verify-md5
  
  # 恢复上传
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --file /path/to/file --upload-id existing-upload-id
  
  # 列举分片
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --upload-id existing-upload-id --action list-parts
  
  # 中止上传
  python s3_multipart_upload.py --bucket my-bucket --key large-file.zip --upload-id existing-upload-id --action abort
        """
    )
    
    # 基本参数
    parser.add_argument('--bucket', required=True, help='桶名称')
    parser.add_argument('--key', required=True, help='对象键名')
    parser.add_argument('--file', help='要上传的文件路径')
    parser.add_argument('--part-size', type=int, default=5*1024*1024, 
                       help='分片大小（字节，默认5MB）')
    parser.add_argument('--verify-md5', action='store_true', 
                       help='启用 Content-MD5 校验')
    parser.add_argument('--upload-id', help='现有的 UploadId（用于恢复上传）')
    parser.add_argument('--action', choices=['upload', 'list-parts', 'abort'], 
                       default='upload', help='操作类型')
    
    # 连接参数
    parser.add_argument('--endpoint', default='http://localhost:8000',
                       help='S3 服务端点 (默认: http://localhost:8000)')
    parser.add_argument('--access-key', default='AKIAIOSFODNN7EXAMPLE',
                       help='AWS Access Key (默认: 教学用凭据)')
    parser.add_argument('--secret-key', default='wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
                       help='AWS Secret Key (默认: 教学用凭据)')
    parser.add_argument('--region', default='us-east-1',
                       help='AWS 区域 (默认: us-east-1)')
    
    # 输出格式
    parser.add_argument('--format', choices=['simple', 'json'], default='simple',
                       help='输出格式 (默认: simple)')
    parser.add_argument('--verbose', action='store_true',
                       help='详细输出')
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()
    
    # 创建客户端
    client = MultipartUploadClient(
        endpoint_url=args.endpoint,
        access_key=args.access_key,
        secret_key=args.secret_key,
        region=args.region
    )
    
    print("🚀 S3 多部分上传工具")
    print("=" * 50)
    print(f"📦 桶: {args.bucket}")
    print(f"🔑 键: {args.key}")
    print(f"🌐 端点: {args.endpoint}")
    print(f"⚙️  操作: {args.action}")
    print("=" * 50)
    
    try:
        if args.action == 'list-parts':
            if not args.upload_id:
                print("❌ 错误: list-parts 操作需要 --upload-id")
                sys.exit(1)
            client.list_parts(args.bucket, args.key, args.upload_id)
            
        elif args.action == 'abort':
            if not args.upload_id:
                print("❌ 错误: abort 操作需要 --upload-id")
                sys.exit(1)
            client.abort_multipart_upload(args.bucket, args.key, args.upload_id)
            
        elif args.action == 'upload':
            if not args.file:
                print("❌ 错误: upload 操作需要 --file")
                sys.exit(1)
            
            success = client.upload_file(
                bucket=args.bucket,
                key=args.key,
                file_path=args.file,
                part_size=args.part_size,
                verify_md5=args.verify_md5,
                upload_id=args.upload_id
            )
            
            if not success:
                sys.exit(1)
        
        print("\n✅ 操作完成！")
        
    except KeyboardInterrupt:
        print("\n⚠️  操作被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 操作失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()